import glob
import pickle
from pathlib import Path

import cv2
import numpy as np

classification = ['airplane',
                  'automobile',
                  'bird',
                  'cat',
                  'deer',
                  'dog',
                  'frog',
                  'horse',
                  'ship',
                  'truck']


def unpickle(file):
    with open(file, 'rb') as fo:
        data = pickle.load(fo, encoding='bytes')
    return data


def get_data_labels(trfiles_):
    data = []
    labels = []
    for file in trfiles_:
        dt = unpickle(file)
        data += list(dt[b"data"])
        labels += list(dt[b"labels"])
    return data, labels


def convert_to_image(data, labels, data_type='train'):
    imgs = np.reshape(data, [-1, 3, 32, 32])
    for i in range(imgs.shape[0]):
        im_data = imgs[i, ...]
        im_data = np.transpose(im_data, [1, 2, 0])
        im_data = cv2.cvtColor(im_data, cv2.COLOR_RGB2BGR)
        if data_type == 'train':
            f = "{}/{}".format("data/image/train", classification[labels[i]])
        else:
            f = "{}/{}".format("data/image/test", classification[labels[i]])
        if not Path(f).exists():
            Path(f).mkdir(parents=True)
        cv2.imwrite("{}/{}.jpg".format(f, str(i)), im_data)


if __name__ == '__main__':
    # Convert to image
    train_files = glob.glob('../data/cifar-10-batches-py/data_batch*')
    data_train, labels_train = get_data_labels(train_files)
    convert_to_image(data_train, labels_train, data_type='train')

    test_files = glob.glob('../data/cifar-10-batches-py/test_batch*')
    data_test, labels_test = get_data_labels(test_files)
    convert_to_image(data_train, labels_train, data_type='train')

    # Convert to tfrecord
